{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# AutoQuant\n",
    "\n",
    "This notebook contains an example of how to use AIMET AutoQuant feature.\n",
    "\n",
    "AIMET offers a suite of neural network post-training quantization (PTQ) techniques that can be applied in succession. However, finding the right sequence of techniques to apply is time-consuming and can be challenging for non-expert users. We instead recommend AutoQuant to save time and effort.\n",
    "\n",
    "AutoQuant is an API that analyzes the model and automatically applies various PTQ techniques based on best-practices heuristics. You specify a tolerable accuracy drop, and AutoQuant applies PTQ techniques cumulatively until the target accuracy is satisfied.\n",
    "\n",
    "## Overall flow\n",
    "\n",
    "This example performs the following steps:\n",
    "\n",
    "1. Instantiate the example evaluation and training pipeline\n",
    "2. Load a pretrained FP32 model\n",
    "3. Determine the baseline FP32 accuracy\n",
    "4. Define constants and helper functions\n",
    "5. Apply AutoQuant\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "\n",
    "Note\n",
    "\n",
    "This notebook does not show state-of-the-art results. For example, it uses a relatively quantization-friendly model (Resnet18). Also, some optimization parameters like number of fine-tuning epochs are chosen to improve execution speed in the notebook.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Dataset\n",
    "\n",
    "This example does image classification on the ImageNet dataset. If you already have a version of the data set, use that. Otherwise download the data set, for example from https://image-net.org/challenges/LSVRC/2012/index .\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "\n",
    "Note\n",
    "\n",
    "To speed up the execution of this notebook, you can use a reduced subset of the ImageNet dataset. For example: The entire ILSVRC2012 dataset has 1000 classes, 1000 training samples per class and 50 validation samples per class. However, for the purpose of running this notebook, you can reduce the dataset to, say, two samples per class.\n",
    "\n",
    "</div>\n",
    "\n",
    "Edit the cell below to specify the directory where the downloaded ImageNet dataset is saved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_DIR = '/path/to/dataset/'         # Replace this path with a real directory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 1. Instantiate the example training and validation pipeline\n",
    "\n",
    "**Use the following training and validation loop for the image classification task.**\n",
    "\n",
    "Things to note:\n",
    "\n",
    "- AIMET does not put limitations on how the training and validation pipeline is written. AIMET modifies the user's model to create a QuantizationSim model, which is still a PyTorch model. The QuantizationSim model can be used in place of the original model when doing inference or training.\n",
    "- AIMET doesn not put limitations on the interface of the `evaluate()` or `train()` methods. You should be able to use your existing evaluate and train routines as-is.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'\n",
    "\n",
    "import tensorflow as tf\n",
    "from aimet_tensorflow.keras.auto_quant import AutoQuant\n",
    "\n",
    "from typing import Optional\n",
    "from Examples.common import image_net_config\n",
    "from Examples.tensorflow.utils.keras.image_net_dataset import ImageNetDataset\n",
    "from Examples.tensorflow.utils.keras.image_net_evaluator import ImageNetEvaluator\n",
    "\n",
    "\n",
    "class ImageNetDataPipeline:\n",
    "    \"\"\"\n",
    "    Provides APIs for model evaluation and finetuning using ImageNet Dataset.\n",
    "    \"\"\"\n",
    "\n",
    "    @staticmethod\n",
    "    def get_val_dataset(batch_size: Optional[int] = None) -> tf.data.Dataset:\n",
    "        \"\"\"\n",
    "        Instantiates a validation dataloader for ImageNet dataset and returns it\n",
    "        :return: A tensorflow dataset\n",
    "        \"\"\"\n",
    "        if batch_size is None:\n",
    "            batch_size = image_net_config.evaluation['batch_size']\n",
    "\n",
    "        data_loader = ImageNetDataset(DATASET_DIR,\n",
    "                                      image_size=image_net_config.dataset['image_size'],\n",
    "                                      batch_size=batch_size)\n",
    "\n",
    "        return data_loader\n",
    "\n",
    "    @staticmethod\n",
    "    def evaluate(model, iterations=None) -> float:\n",
    "        \"\"\"\n",
    "        Given a Keras model, evaluates its Top-1 accuracy on the validation dataset\n",
    "        :param model: The Keras model to be evaluated.\n",
    "        :param iterations: The number of iterations to run. If None, all the data will be used\n",
    "        :return: The accuracy for the sample with the maximum accuracy.\n",
    "        \"\"\"\n",
    "        evaluator = ImageNetEvaluator(DATASET_DIR,\n",
    "                                      image_size=image_net_config.dataset[\"image_size\"],\n",
    "                                      batch_size=image_net_config.evaluation[\"batch_size\"])\n",
    "\n",
    "        return evaluator.evaluate(model=model, iterations=iterations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Load a pretrained FP32 model\n",
    "\n",
    "**Load a pretrained resnet50 model from Keras.** \n",
    "\n",
    "You can load any pretrained PyTorch model instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from tensorflow.keras.applications.resnet import ResNet50\n",
    "\n",
    "model = ResNet50(weights='imagenet')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## 3.  Determine the baseline FP32 accuracy\n",
    "\n",
    "**Determine the floating point 32-bit (FP32) accuracy of this model using the `evaluate()` routine.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ImageNetDataPipeline.evaluate(model=model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Define constants and helper functions\n",
    "\n",
    "**4.1 Define the following constants:**\n",
    "\n",
    "- **EVAL_DATASET_SIZE** A typical value is 5000. In this example, the value has been set to 50 for faster execution.\n",
    "- **CALIBRATION_DATASET_SIZE** A typical value is 2000. In this example, the value has been set to 20 for faster execution.\n",
    "- **BATCH_SIZE** You define the batch size. Set to 10 in this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "EVAL_DATASET_SIZE = 50\n",
    "CALIBRATION_DATASET_SIZE = 20\n",
    "BATCH_SIZE = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**4.2 Use the constants to create the evaluation dataset.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "eval_dataset = ImageNetDataPipeline.get_val_dataset(BATCH_SIZE).dataset\n",
    "unlabeled_dataset = eval_dataset.map(lambda images, labels: images)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "**4.3 Prepare the evaluation callback function.**\n",
    "\n",
    "The **eval_callback()** function takes the model object to evaluate and compile option dictionary and the number of samples to use as arguments. If the **num_samples** argument is None, the whole evaluation dataset is used to evaluate the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from typing import Optional\n",
    "\n",
    "\n",
    "def eval_callback(model: tf.keras.Model,\n",
    "                  num_samples: Optional[int] = None) -> float:\n",
    "    if num_samples is None:\n",
    "        num_samples = EVAL_DATASET_SIZE\n",
    "\n",
    "    sampled_dataset = eval_dataset.take(num_samples)\n",
    "\n",
    "    # Model should be compiled before evaluation\n",
    "    model.compile(optimizer=tf.keras.optimizers.Adam(),\n",
    "                  loss=tf.keras.losses.CategoricalCrossentropy(),\n",
    "                  metrics=tf.keras.metrics.CategoricalAccuracy())\n",
    "    _, acc = model.evaluate(sampled_dataset)\n",
    "\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Apply AutoQuant\n",
    "\n",
    "**5.1 Create an AutoQuant object.**\n",
    "\n",
    "The AutoQuant feature uses an unlabeled dataset to quantize the model. The **UnlabeledDatasetWrapper** class creates an unlabeled Dataset object from a labeled Dataset.\n",
    "\n",
    "The **allowed_accuracy_drop** indicates the tolerable accuracy drop. AutoQuant applies a series of quantization features until the target accuracy (FP32 accuracy - allowed accuracy drop) is satisfied. When the target accuracy is reached, AutoQuant returns immediately without applying furhter PTQ techniques. See the [AutoQuant User Guide](https://quic.github.io/aimet-pages/releases/latest/user_guide/auto_quant.html) and [AutoQuant API documentation](https://quic.github.io/aimet-pages/releases/latest/api_docs/torch_auto_quant.html) for details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "auto_quant = AutoQuant(allowed_accuracy_drop=0.01,\n",
    "                       unlabeled_dataset=unlabeled_dataset,\n",
    "                       eval_callback=eval_callback)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**5.2 Set AdaRound Parameters (optional)**\n",
    "\n",
    "AutoQuant uses predefined default parameters for AdaRound.\n",
    "These values were determined empirically and work well with the common models.\n",
    "\n",
    "If necessary, you can use custom parameters for Adaround.\n",
    "This example uses very small AdaRound parameters for faster execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from aimet_tensorflow.adaround.adaround_weight import AdaroundParameters\n",
    "\n",
    "ADAROUND_DATASET_SIZE = 2000\n",
    "adaround_dataset = unlabeled_dataset.take(ADAROUND_DATASET_SIZE)\n",
    "adaround_params = AdaroundParameters(adaround_dataset,\n",
    "                                     num_batches=ADAROUND_DATASET_SIZE // BATCH_SIZE)\n",
    "auto_quant.set_adaround_params(adaround_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**5.3 Run AutoQuant Optimization**\n",
    "\n",
    "This step runs AutoQuant optimization. AutoQuant returns the following:\n",
    "- The best possible quantized model\n",
    "- The corresponding evaluation score\n",
    "- The path to the encoding file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "model, accuracy, encoding_path = auto_quant.apply(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## For more information\n",
    "\n",
    "See the [AIMET API docs](https://quic.github.io/aimet-pages/AimetDocs/api_docs/index.html) for details about the AIMET APIs and optional parameters.\n",
    "\n",
    "See the [other example notebooks](https://github.com/quic/aimet/tree/develop/Examples/torch/quantization) to learn how to use other AIMET post-training quantization techniques."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
